# -*- coding: utf-8 -*-
import pandas as pd
import requests
import json
import time
# import matplotlib.pyplot as plt
import os
import argparse
import threading
import base64
import queue
from Crypto.Hash import HMAC
from Crypto.Hash import SHA1
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5
# 创建一个队列
shared_queue = queue.Queue()

# 保存测试结果
class ResponseResult():
    id: str
    status: str
    question: str
    response: str
    first_time: float
    total_time: float
    tokens_num: int

url_chat = '0.0.0.0:6000'
image_dataset_path = r""
# def send_request_and_resolve(send_data):
#     json_data = json.dumps(send_data)
#
#     header_dict = {'Content-Type': 'application/json'}
#
#     r = requests.post(url_chat, headers=header_dict, data=json_data, timeout=300)
#
#     return r.content

# 应用的accessKeyId和accessKeySecret
accesskey = "3024d9f88e6e160caeddb810d0bab03d"
accessKeySecret = "7fa134b45034c08bfa2f18073b8f5c60"

# 请求时间
timestamp = int(time.time() * 1000)

# 加密类型：SHA1
signtype = "SHA1"

# 入参（这里假设请求体是JSON格式的字符串）
params = '{"key1": "value1", "key2": "value2"}'  # 根据您的实际请求体内容修改

# base64加密
appPolicy = accesskey + accessKeySecret + str(timestamp) + params[:512]
appPolicy = appPolicy.encode('utf-8')

# HmacSHA1加密
hmac = HMAC.new(accessKeySecret.encode('utf-8'), digestmod=SHA1)
hmac.update(appPolicy)
sign = base64.b64encode(hmac.digest()).decode('utf-8')


def get_response(id, prompt, image_name):
    rep_res = ResponseResult()
    image_url = image_dataset_path + image_name
    with open(image_dataset_path + "\\" + image_name, 'rb') as f:
        image = str(base64.b64encode(f.read()), 'utf-8')
    print("image_url: " + image)
    # # 设置请求头
    header_dict = {'Content-Type': 'application/json'}
    # header_dict = {
    #     "Accesskey": accesskey,
    #     "Signtype": signtype,
    #     "Timestamp": str(timestamp),
    #     "Sign": sign,
    #     "Content-Type": "application/json"
    # }
    # openai
    # send_data = {
    #     "model": "string",
    #     "messages": [
    #         {
    #             "role": "user",
    #             "content": [
    #                 {
    #                     "type": "image_url",
    #                     "image_url": image
    #                 },
    #                 {
    #                     "type": "text",
    #                     "text": prompt
    #                 }
    #             ]
    #         }
    #     ],
    #     "max_tokens": 512,
    #     "presence_penalty": 1.03,
    #     "frequency_penalty": 1.0,
    #     "temperature": 0.5,
    #     "top_p": 0.96,
    #     "stream": False,
    #     "stop": [
    #         "lf"
    #     ]
    # }

    # send_data = {
    #     "inputs":[
    #         {
    #             "type": "text",
    #             "text": "A chat between a curious user and an artificial intelligence assistant.The assistant gives helpful,detailed,and polite answers to the user questions. USER: <image> {}. ASSISTANT:".format(prompt)
    #         },
    #         {
    #             "type": "image_url",
    #             "image_url": image
    #         }
    #     ],
    #     "parameters": {
    #         "do_sample": False,
    #         "max_new_tokens": 256
    #     },
    #     "stream": True
    # }

    send_data = {
        "model": "string",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "Describe the content in the image."
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": image
                        }
                    }
                ]
            }
        ],
        "max_tokens": 300
    }


    # send_data = {
    #     "model": "string",
    #     "messages": [
    #         {
    #             "role": "user",
    #             "content": "<img></img>\n请问这是哪个城市的建筑物？",
    #             "function_call": {}
    #         }
    #     ],
    #     "stop": [],
    #     "stream": False
    # }

    json_data = json.dumps(send_data)


    start_time = time.monotonic()
    first_time = -1
    end_time = 0
    rep_lst = []
    try:
        # 流式获取模型输出，需自行处理
        with requests.post(url_chat, headers=header_dict, data=json_data, timeout=300, stream=True) as response:
            print("response.text:  " + response.text)
            for chunk in response.iter_lines(chunk_size=None):
                chunk = chunk.strip()
                if not chunk:
                    continue
                if first_time<0:
                    first_time = time.monotonic() - start_time
                rep_lst.append(chunk)
            end_time = time.monotonic() - start_time
            chunk = rep_lst[-1].decode('utf-8')
            chunk = chunk[5:]
            data = json.loads(chunk)['text'].decode('utf-8')
    
        # 生成测试结果
        rep_res.id = id
        rep_res.question = prompt
        rep_res.response = data
        rep_res.first_time = first_time
        rep_res.total_time = end_time
        rep_res.tokens_num = len(rep_lst)
        rep_res.status = "success"
    except:
        
        rep_res.id = id
        rep_res.question = prompt
        rep_res.response = "\n".join(rep_lst)
        rep_res.status = "fail"
    return rep_res

def gen_res(id, test_time, image_list):
    res_lst = []
    start_time = time.monotonic()
    for i in range(1000000):
        idx = i % len(image_list)
        image = image_list[idx]
        prompt = "describe the image"
        new_id = str(id) + '-' + str(i)
        res_lst.append(get_response(new_id, prompt, image))
        if time.monotonic()-start_time > test_time:
            break
    shared_queue.put(res_lst)


if __name__ == '__main__':

    thread_num = 1  # 并发线程数
    test_time = 30 # 测试时间，单位是秒
    # # 启用命令行传参
    # parser = argparse.ArgumentParser()
    # parser.add_argument('--thread_num', type=str, help='并发线程数')
    # parser.add_argument('--test_time', type=int, help='测试时间，单位是秒')
    # args = parser.parse_args()
    # thread_num = args.thread_num
    # test_time = args.test_time

    #df_test_data = pd.read_excel('data/runse-data.xlsx')  # 测试数据

    image_list = os.listdir(image_dataset_path)
    # 创建线程列表
    threads = []
    # 创建和启动线程
    for i in range(thread_num):
        thread = threading.Thread(target=gen_res, args=(i, test_time, image_list))
        threads.append(thread)
        thread.start()
    # 等待所有线程完成
    for thread in threads:
        thread.join()
    print("所有请求已完成。")

    df_res = {
        'id': [],
        'question': [],
        'response': [],
        'first_time': [],
        'total_time': [],
        'tokens_num': []
    }
    fail_dict = {
        'id': [],
        'question': [],
    }

    for i in range(shared_queue.qsize()):
        # 从队列中获取信息
        res_lst = shared_queue.get(timeout=1000)
        for res in res_lst:
            if res.status == "success":
                df_res['id'].append(res.id)
                df_res['question'].append(res.question)
                df_res['response'].append(res.response)
                df_res['first_time'].append(res.first_time)
                df_res['total_time'].append(res.total_time)
                df_res['tokens_num'].append(res.tokens_num)
            else:
                fail_dict['id'].append(res.id)
                fail_dict['question'].append(res.question)
    print(f"成功请求数量：{len(df_res['id'])}  失败请求数量：{len(fail_dict['id'])}")
    ticks = time.strftime('%Y-%m-%d_%H-%M-%S',time.localtime())
    folder_path = f"result/runse/{ticks}_tn-{str(thread_num)}_tt-{str(test_time)}"
    # 检查文件夹是否存在
    if not os.path.exists(folder_path):
        # 如果文件夹不存在，创建它
        os.makedirs(folder_path)
        print(f"文件夹 '{folder_path}' 已创建。")
    else:
        print(f"文件夹 '{folder_path}' 已存在。")
    df_res = pd.DataFrame.from_dict(df_res)
    df_res.to_excel(f'{folder_path}/result.xlsx', index=False)
    fail_dict = pd.DataFrame.from_dict(fail_dict)
    fail_dict.to_excel(f'{folder_path}/fail.xlsx', index=False)
    first_time_avg = sum(df_res['first_time']) / len(df_res['first_time'])
    # req_num = len(df_res['first_time'])
    # 保存结果
    len_total = len(df_res['id'])+len(fail_dict['id'])
    acc = round(len(df_res['id'])/len_total*100)
    token_speed = 0
    prompt_len = 0
    output_len = 0
    context_len = 0
    for i in range(df_res.shape[0]):
        token_speed += df_res.loc[i, "tokens_num"] /  df_res.loc[i, "total_time"]
    token_speed /= df_res.shape[0]
    res_info = f'''测试线程数：{thread_num}
测试时间：{test_time}秒
请求数量：{len_total}  成功请求数量：{len(df_res['id'])}  失败请求数量：{len(fail_dict['id'])}
事务成功率：{acc}%
平均首字响应时间：{first_time_avg}
每秒输出token：{token_speed}
'''
    print(res_info)
    with open(f'{folder_path}/log', 'w', encoding='utf8') as f:
        f.write(res_info)
    
    # # 画图
    # x = [len(s) for s in df_res['question']]
    # y = [x for x in df_res['first_time']]
    # min_y = min(y)
    # max_y = max(y)
    # plt.ylim(0, max_y+min_y)
    # plt.plot(x, y, 'ro')
    # plt.ylabel('first token response time')
    # plt.xlabel('prompt length')
    # plt.savefig(f'{folder_path}/ttft.jpg')
    # # plt.show()
    